package wrr

import (
	"context"
	"fmt"
	"google.golang.org/grpc"
	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/balancer/base"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/status"
	"io"
	"sync"
)

const WeightRoundBin = "custom_wrr"

func newBuilder() balancer.Builder {
	// NewBalancerBuilder 是帮吧picker builder转化为balancer.Builder
	return base.NewBalancerBuilder(WeightRoundBin, &PickerBuilder{}, base.Config{HealthCheck: true})
}
func init() {
	balancer.Register(newBuilder())
}

type PickerBuilder struct {
}

func (p *PickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
	conns := make([]*weightConn, 0, len(info.ReadySCs))
	for subConn, connInfo := range info.ReadySCs {
		weightVal, _ := connInfo.Address.Metadata.(map[string]any)
		weight, _ := weightVal["weight"].(float64)
		mp, _ := weightVal["server"].(map[string]interface{})
		fmt.Println(mp)
		conns = append(conns, &weightConn{
			SubConn:       subConn,
			weight:        int(weight),
			currentWeight: int(weight),
			//server:        server,
			available: true,
		})
	}
	return &Picker{
		conns: conns,
		mutex: &sync.Mutex{},
	}
}

type Picker struct {
	//真正执行负载均衡的地方
	conns []*weightConn
	mutex *sync.Mutex
}

func (p *Picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
	if len(p.conns) == 0 {
		return balancer.PickResult{}, nil
	}

	p.mutex.Lock()
	defer p.mutex.Unlock()
	var total int
	var maxCC *weightConn
	for _, conn := range p.conns {
		if !conn.available {
			continue
		}

		total += conn.weight
		conn.currentWeight += conn.weight
		if maxCC == nil || conn.currentWeight > maxCC.currentWeight {
			maxCC = conn
		}
	}
	//最大权重
	maxCC.currentWeight -= total
	return balancer.PickResult{
		SubConn: maxCC.SubConn,
		Done: func(info balancer.DoneInfo) {
			err := info.Err
			if err == nil {
				return
			}
			switch err {
			case context.Canceled:
				return
			case io.EOF:
				//基本不可用
				maxCC.available = false
			case context.DeadlineExceeded:
			default:
				st, ok := status.FromError(err)
				if ok {
					code := st.Code()
					switch code {
					case codes.Unavailable:
						maxCC.available = false
						go func() {
							//额外的goroutine去探活
							//if p.healthCheck(maxCC) {
							//	maxCC.available = true
							//	//最好加点流量控制的措施
							//}
						}()
					case codes.ResourceExhausted:
						//限流，降低权重
						maxCC.currentWeight -= total
						maxCC.weight -= 3
					}
				}
			}
		},
	}, nil
}
func (p *Picker) healthCheck(wc *weightConn) bool {
	//注册服务健康检查
	grpc_health_v1.RegisterHealthServer(wc.server, health.NewServer())
	return true
}

type weightConn struct {
	//初始权重
	weight int
	//当前权重
	currentWeight int
	balancer.SubConn
	available bool
	server    *grpc.Server
}
